/*
 *  Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
 *  <p>
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  <p>
 *  http://www.apache.org/licenses/LICENSE-2.0
 *  <p>
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package com.agentsflex.image.tencent;

import com.agentsflex.core.image.*;
import com.agentsflex.core.llm.client.HttpClient;
import com.agentsflex.core.util.Maps;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.tencentcloudapi.common.DatatypeConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.text.SimpleDateFormat;
import java.util.*;

public class TencentImageModel implements ImageModel {
    private static final Logger LOG = LoggerFactory.getLogger(TencentImageModel.class);
    private final TencentImageModelConfig config;
    private final HttpClient httpClient = new HttpClient();

    public TencentImageModel(TencentImageModelConfig config) {
        this.config = config;
    }

    @Override
    public ImageResponse generate(GenerateImageRequest request) {
        try {
            String payload = promptToPayload(request);
            Map<String, String> headers = createAuthorizationToken("SubmitHunyuanImageJob", payload);
            String response = httpClient.post(config.getEndpoint(), headers, payload);
            JSONObject jsonObject = JSON.parseObject(response);
            JSONObject error = jsonObject.getJSONObject("Response").getJSONObject("Error");
            if (error != null && !error.isEmpty()) {
                return ImageResponse.error(error.getString("Message"));
            }
            Object jobId = jsonObject.getJSONObject("Response").get("JobId");
            if (Objects.isNull(jobId)) {
                return ImageResponse.error("response is no jobId");
            }
            String id = (String) jobId;
            return getImage(id);
        } catch (Exception e) {
            return ImageResponse.error(e.getMessage());
        }
    }

    @Override
    public ImageResponse img2imggenerate(GenerateImageRequest request) {
        return null;
    }

    @Override
    public ImageResponse edit(EditImageRequest request) {
        throw new IllegalStateException("TencentImageModel Can not support edit image.");
    }

    @Override
    public ImageResponse vary(VaryImageRequest request) {
        throw new IllegalStateException("TencentImageModel Can not support vary image.");
    }


    private static final Object LOCK = new Object();

    private ImageResponse getImage(String jobId) {
        ImageResponse imageResponse = null;
        while (true) {
            synchronized (LOCK) {
                imageResponse = callService(jobId);
                if (!Objects.isNull(imageResponse)) {
                    break;
                }
                // 等待一段时间再重试
                try {
                    LOCK.wait(1000);
                } catch (InterruptedException e) {
                    // 线程在等待时被中断
                    Thread.currentThread().interrupt();
                    imageResponse = ImageResponse.error(e.toString());
                    break;
                }
            }
        }
        return imageResponse;
    }


    public ImageResponse callService(String jobId) {
        try {
            String payload = Maps.of("JobId", jobId).toJSON();
            Map<String, String> headers = createAuthorizationToken("QueryHunyuanImageJob", payload);
            String resp = httpClient.post(config.getEndpoint(), headers, payload);
            JSONObject resultJson = JSONObject.parseObject(resp).getJSONObject("Response");
            JSONObject error = resultJson.getJSONObject("Error");
            if (error != null && !error.isEmpty()) {
                return ImageResponse.error(error.getString("Message"));
            }
            if (Objects.isNull(resultJson.get("JobStatusCode"))) {
                return ImageResponse.error("response is no JobStatusCode");
            }
            Integer jobStatusCode = resultJson.getInteger("JobStatusCode");
            if (Objects.equals(5, jobStatusCode)) {
                //处理完成
                if (Objects.isNull(resultJson.get("ResultImage"))) {
                    return ImageResponse.error("response is no ResultImage");
                }
                JSONArray imagesArray = resultJson.getJSONArray("ResultImage");
                ImageResponse response = new ImageResponse();
                for (int i = 0; i < imagesArray.size(); i++) {
                    String imageObj = imagesArray.getString(i);
                    response.addImage(imageObj);
                }
                return response;
            }
            if (Objects.equals(4, jobStatusCode)) {
                //处理错误
                return ImageResponse.error(resultJson.getString("JobErrorMsg"));
            }
        } catch (Exception e) {
            return ImageResponse.error(e.getMessage());
        }
        return null;
    }


    public static String promptToPayload(GenerateImageRequest request) {
        return Maps.of("Prompt", request.getPrompt())
            .setIfNotEmpty("NegativePrompt", request.getNegativePrompt())
            .setIfNotEmpty("Style", request.getSize())
            .setIfNotEmpty("Resolution", request.getQuality())
            .setIfNotEmpty("Num", request.getN())
            .setIfNotEmpty(request.getOptions())
            .toJSON();
    }


    private final static Charset UTF8 = StandardCharsets.UTF_8;
    private final static String CT_JSON = "application/json; charset=utf-8";

    public static byte[] hmac256(byte[] key, String msg) throws Exception {
        Mac mac = Mac.getInstance("HmacSHA256");
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, mac.getAlgorithm());
        mac.init(secretKeySpec);
        return mac.doFinal(msg.getBytes(UTF8));
    }

    public static String sha256Hex(String s) throws Exception {
        MessageDigest md = MessageDigest.getInstance("SHA-256");
        byte[] d = md.digest(s.getBytes(UTF8));
        return DatatypeConverter.printHexBinary(d).toLowerCase();
    }

    /**
     * @return java.util.Map<java.lang.String, java.lang.String>
     * @Author sunch
     * @Description 封装参数
     * @Date 17:34 2025/3/5
     * @Param [action, payload]
     */
    public Map<String, String> createAuthorizationToken(String action, String payload) {
        try {
            String service = config.getService();
            String host = config.getHost();
            String version = "2023-09-01";
            String algorithm = "TC3-HMAC-SHA256";
            String timestamp = String.valueOf(System.currentTimeMillis() / 1000);
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd");
            // 注意时区，否则容易出错
            sdf.setTimeZone(TimeZone.getTimeZone("UTC"));
            String date = sdf.format(new Date(Long.parseLong(timestamp + "000")));

            // ************* 步骤 1：拼接规范请求串 *************
            String httpRequestMethod = "POST";
            String canonicalUri = "/";
            String canonicalQueryString = "";
            String canonicalHeaders = "content-type:application/json; charset=utf-8\n"
                + "host:" + host + "\n" + "x-tc-action:" + action.toLowerCase() + "\n";
            String signedHeaders = "content-type;host;x-tc-action";

            String hashedRequestPayload = sha256Hex(payload);
            String canonicalRequest = httpRequestMethod + "\n" + canonicalUri + "\n" + canonicalQueryString + "\n"
                + canonicalHeaders + "\n" + signedHeaders + "\n" + hashedRequestPayload;
//            System.out.println(canonicalRequest);

            // ************* 步骤 2：拼接待签名字符串 *************
            String credentialScope = date + "/" + service + "/" + "tc3_request";
            String hashedCanonicalRequest = sha256Hex(canonicalRequest);
            String stringToSign = algorithm + "\n" + timestamp + "\n" + credentialScope + "\n" + hashedCanonicalRequest;
//            System.out.println(stringToSign);

            // ************* 步骤 3：计算签名 *************
            byte[] secretDate = hmac256(("TC3" + config.getApiKey()).getBytes(UTF8), date);
            byte[] secretService = hmac256(secretDate, service);
            byte[] secretSigning = hmac256(secretService, "tc3_request");
            String signature = DatatypeConverter.printHexBinary(hmac256(secretSigning, stringToSign)).toLowerCase();
//            System.out.println(signature);

            // ************* 步骤 4：拼接 Authorization *************
            String authorization = algorithm + " " + "Credential=" + config.getApiSecret() + "/" + credentialScope + ", "
                + "SignedHeaders=" + signedHeaders + ", " + "Signature=" + signature;
//            System.out.println(authorization);

            Map<String, String> headers = new HashMap<>();
            headers.put("Authorization", authorization);
            headers.put("Content-Type", CT_JSON);
            headers.put("Host", host);
            headers.put("X-TC-Action", action);
            headers.put("X-TC-Timestamp", timestamp);
            headers.put("X-TC-Version", version);
            headers.put("X-TC-Region", config.getRegion());
            return headers;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
